from fastmcp import Client
import asyncio
import logging
import re

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger("calculator_mcp_client")

class CalculatorClient:
    def __init__(self, host="localhost", port=25565):
        # 根据server.py中的配置设置默认连接参数
        self.host = host
        self.port = port
        self.client = None
        self.running = False

    async def start(self):
        try:
            # 构建服务器连接URL，与server.py中的HTTP传输协议匹配
            transport_url = f"http://{self.host}:{self.port}/mcp"
            self.client = Client(transport=transport_url)
            self.running = True
            
            print(f"正在连接到计算服务器 {self.host}:{self.port}...")
            
            # 使用异步上下文管理器确保正确连接和断开
            async with self.client:
                print("已连接到计算服务器！")
                print("可用工具: add, subtract, multiply, divide")
                print("输入函数调用（如 add(2, 3)），输入 exit 退出")
                
                # 启动输入循环
                await self.input_handler()
        except Exception as e:
            logger.error(f"连接服务器失败: {e}")
            print(f"连接错误: {e}")
        finally:
            await self.cleanup()
            print("已断开连接")

    async def call_tool(self, tool_name: str, *args):
        """调用服务器端的工具函数，使用fastmcp的API"""
        if not self.client or not self.running:
            raise RuntimeError("客户端未连接或已停止")
            
        # 构建工具调用参数，与server.py中的工具参数匹配
        if len(args) >= 2:
            arguments = {"a": args[0], "b": args[1]}
        else:
            raise ValueError("需要2个参数")
            
        logger.info(f"调用工具: {tool_name}，参数: {arguments}")
        
        # 使用fastmcp的call_tool方法调用服务端工具
        result = await self.client.call_tool(
            tool_name,
            arguments=arguments
        )

        return result

    async def input_handler(self):
        """处理用户输入，支持函数调用格式"""
        # 用于解析函数调用的正则表达式
        function_pattern = re.compile(r'^(\w+)\(([^)]*)\)$')
        
        while self.running:
            cmd = await asyncio.to_thread(input, "输入命令: ")
            if cmd.lower() == "exit":
                self.running = False
                break
            
            # 尝试解析函数调用格式（如 add(2, 3)）
            match = function_pattern.match(cmd)
            if match:
                function_name = match.group(1)
                args_str = match.group(2)
                
                try:
                    # 转换参数为数字
                    args = [float(arg.strip()) for arg in args_str.split(',') if arg.strip()]
                    
                    # 调用服务端的工具函数
                    result = await self.call_tool(function_name, *args)
                    print(f"计算结果: {result}")
                except ValueError as e:
                    print(f"参数错误: {e}")
                except Exception as e:
                    print(f"调用失败: {e}")
            else:
                print("无效的输入格式，请使用函数调用格式，例如: add(2, 3)")

    async def cleanup(self):
        """清理资源"""
        self.running = False

# 主函数
async def main():
    # 获取连接参数，默认连接到localhost:25565（对应server.py中的0.0.0.0:25565）
    host = input("服务器地址 (默认: localhost): ") or "localhost"
    port_input = input("端口 (默认: 25565): ")
    port = int(port_input) if port_input else 25565
    
    # 启动客户端
    client = CalculatorClient(host, port)
    await client.start()

if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("客户端已停止")